import torch
import torch.nn as nn
from torch.nn import functional as F

class Expert(nn.Module):
    """专家网络类定义"""
    def __init__(self):
        super(Expert, self).__init__()
        # 特征提取层
        self.features = nn.Sequential(
            # 第一个卷积块：1->32通道
            nn.Conv2d(1, 32, 3, padding=1),  # 保持图像大小不变
            nn.BatchNorm2d(32),  # 批归一化提高训练稳定性
            nn.ReLU(),  # 激活函数
            
            # 第二个卷积块：32->64通道
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),  # 图像尺寸减半
            nn.Dropout2d(0.25),  # 防止过拟合
            
            # 第三个卷积块：64->128通道
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2),  # 再次减半
            nn.Dropout2d(0.25),
        )
        
        # 分类层
        self.classifier = nn.Sequential(
            nn.Linear(128 * 7 * 7, 512),  # 展平后的特征图连接到全连接层
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(0.5),  # 更强的dropout防止过拟合
            nn.Linear(512, 10),  # 最终输出10个类别
        )

    def forward(self, x):
        """前向传播
        
        参数:
            x: 输入图像 [batch_size, 1, 28, 28]
        返回:
            输出: 类别预测 [batch_size, 10]
        """
        x = self.features(x)  # 提取特征
        x = x.view(x.size(0), -1)  # 展平特征图
        return self.classifier(x)  # 分类预测


class Gate(nn.Module):
    """门控网络类定义"""
    def __init__(self, num_experts):
        super(Gate, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        
        self.attention = nn.Sequential(
            nn.Linear(32 * 7 * 7, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, num_experts),
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return F.softmax(self.attention(x) / 0.5, dim=1)


class MOE_Model(nn.Module):
    """混合专家模型类定义"""
    def __init__(self, num_experts):
        super(MOE_Model, self).__init__()
        self.num_experts = num_experts
        self.experts = nn.ModuleList([Expert() for _ in range(num_experts)])
        self.gate = Gate(num_experts)

    def forward(self, x):
        gate_output = self.gate(x)
        expert_outputs = torch.stack([expert(x) for expert in self.experts], dim=1)
        return (expert_outputs * gate_output.unsqueeze(-1)).sum(dim=1)


class EarlyStopping:
    """早停机制类定义"""
    def __init__(self, patience=7, min_delta=1e-4, mode="max"):
        self.patience = patience
        self.min_delta = min_delta
        self.mode = mode
        self.counter = 0
        self.best_value = None
        self.should_stop = False
        self.best_epoch = 0

    def __call__(self, epoch, value):
        """检查是否应该停止训练
        
        参数:
            epoch (int): 当前训练轮数
            value (float): 当前监控的指标值
            
        返回:
            bool: 是否应该停止训练
        """
        if self.best_value is None:
            self.best_value = value
            self.best_epoch = epoch
            return False

        if self.mode == "min":
            improved = (self.best_value - value) > self.min_delta
        else:
            improved = (value - self.best_value) > self.min_delta

        if improved:
            self.best_value = value
            self.best_epoch = epoch
            self.counter = 0
        else:
            self.counter += 1

        if self.counter >= self.patience:
            self.should_stop = True

        return self.should_stop 